"""
SQLite审计数据批量写入器
优化的终端审计数据写入性能
"""

import sqlite3
import threading
import queue
import time
import logging
from django.conf import settings
from django.db import connection

logger = logging.getLogger(__name__)


class SQLiteAuditWriter:
    """SQLite审计数据批量写入器"""
    
    def __init__(self, batch_size=100, flush_interval=5):
        self.batch_size = batch_size
        self.flush_interval = flush_interval
        self.event_queue = queue.Queue()
        self.command_queue = queue.Queue()
        self.running = True
        
        # 启动后台写入线程
        self.writer_thread = threading.Thread(target=self._batch_writer, daemon=True)
        self.writer_thread.start()
        logger.info("SQLite审计写入器已启动")
        
    def _get_connection(self):
        """获取优化的SQLite连接"""
        db_path = settings.DATABASES['default']['NAME']
        conn = sqlite3.connect(db_path, timeout=30.0)
        
        # 应用性能优化设置
        conn.execute('PRAGMA journal_mode=WAL')
        conn.execute('PRAGMA synchronous=NORMAL')
        conn.execute('PRAGMA cache_size=10000')
        conn.execute('PRAGMA temp_store=MEMORY')
        conn.execute('PRAGMA mmap_size=268435456')
        
        return conn
        
    def add_event(self, session_id, timestamp, event_type, data):
        """添加终端事件到队列"""
        try:
            self.event_queue.put({
                'session_id': session_id,
                'timestamp': timestamp,
                'event_type': event_type,
                'data': data
            }, timeout=1.0)
        except queue.Full:
            logger.warning("事件队列已满，丢弃事件")
        
    def add_command(self, session_id, command, timestamp, exit_code=None, duration=None):
        """添加命令到队列"""
        try:
            self.command_queue.put({
                'session_id': session_id,
                'command': command,
                'timestamp': timestamp,
                'exit_code': exit_code,
                'duration': duration
            }, timeout=1.0)
        except queue.Full:
            logger.warning("命令队列已满，丢弃命令")
        
    def _batch_writer(self):
        """后台批量写入线程"""
        events_batch = []
        commands_batch = []
        last_flush = time.time()
        
        while self.running:
            try:
                # 收集事件数据
                while len(events_batch) < self.batch_size:
                    try:
                        event = self.event_queue.get(timeout=0.1)
                        events_batch.append(event)
                    except queue.Empty:
                        break
                        
                # 收集命令数据
                while len(commands_batch) < self.batch_size:
                    try:
                        command = self.command_queue.get(timeout=0.1)
                        commands_batch.append(command)
                    except queue.Empty:
                        break
                
                # 检查是否需要刷新
                now = time.time()
                should_flush = (
                    len(events_batch) >= self.batch_size or
                    len(commands_batch) >= self.batch_size or
                    (now - last_flush) >= self.flush_interval
                )
                
                if should_flush and (events_batch or commands_batch):
                    self._flush_batches(events_batch, commands_batch)
                    events_batch.clear()
                    commands_batch.clear()
                    last_flush = now
                    
            except Exception as e:
                logger.error(f"批量写入错误: {e}")
                time.sleep(1)
                
    def _flush_batches(self, events_batch, commands_batch):
        """刷新批次数据到数据库"""
        conn = self._get_connection()
        try:
            with conn:
                # 批量插入事件
                if events_batch:
                    conn.executemany('''
                        INSERT INTO cmdb_terminal_event (session_id, timestamp, event_type, data)
                        VALUES (?, ?, ?, ?)
                    ''', [(e['session_id'], e['timestamp'], e['event_type'], e['data']) 
                          for e in events_batch])
                    logger.debug(f"批量插入 {len(events_batch)} 个事件")
                
                # 批量插入命令
                if commands_batch:
                    conn.executemany('''
                        INSERT INTO cmdb_terminal_command (session_id, command, timestamp, exit_code, duration)
                        VALUES (?, ?, ?, ?, ?)
                    ''', [(c['session_id'], c['command'], c['timestamp'], c['exit_code'], c['duration'])
                          for c in commands_batch])
                    logger.debug(f"批量插入 {len(commands_batch)} 个命令")
                          
        except Exception as e:
            logger.error(f"数据库写入错误: {e}")
        finally:
            conn.close()
            
    def stop(self):
        """停止写入器"""
        self.running = False
        if self.writer_thread.is_alive():
            self.writer_thread.join(timeout=5)
        logger.info("SQLite审计写入器已停止")

# 全局审计写入器实例
audit_writer = None

def get_audit_writer():
    """获取全局审计写入器实例"""
    global audit_writer
    if audit_writer is None:
        audit_writer = SQLiteAuditWriter()
    return audit_writer

def stop_audit_writer():
    """停止全局审计写入器"""
    global audit_writer
    if audit_writer:
        audit_writer.stop()
        audit_writer = None
